package typechecker;


import java.util.Arrays;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

import typechecker.ast.Context;
import typechecker.ast.term.CaseAnalisys;
import typechecker.ast.term.DataConstructor;
import typechecker.ast.term.Expression;
import typechecker.ast.term.LambdaAbstraction;
import typechecker.ast.term.LambdaApplication;
import typechecker.ast.term.Letrec;
import typechecker.ast.term.TermVariable;
import typechecker.ast.term.TypeAbstraction;
import typechecker.ast.term.TypeApplication;
import typechecker.ast.type.ForallType;
import typechecker.ast.type.FunctionType;
import typechecker.ast.type.TypeConstructorApplication;
import typechecker.ast.type.TypeExpression;
import typechecker.ast.type.TypeVariable;
import typechecker.walkers.InstConstructorType;
import typechecker.walkers.NTAnnot;
import typechecker.walkers.SubtypingConstraintGenerator;
import typechecker.walkers.TypeAnnot;
import typechecker.walkers.TypeSubst;
import typechecker.walkers.TypeSucc;

public class InferType {	
	public static int errors = 0;
	
	public TypeExpression inferVariable(TermVariable var, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		return Gamma.get(var);
	}

	public TypeExpression inferLambdaAbstraction(LambdaAbstraction labs, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		Expression body = labs.getBody();
		TypeExpression type = labs.getType();
		TermVariable var = labs.getVar();
			
		annot(v, type);
		Gamma.put( var, type);
		TypeExpression t2 = infer(body, c, v, Gamma, ctx);
		Gamma.remove(var);

		return new FunctionType(type, t2);
	}

	public TypeExpression inferTypeAbstraction(TypeAbstraction tabs, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		Expression body = tabs.getBody();
		TypeVariable var = tabs.getTvar();

		ctx.push(var);
		TypeExpression bodyType = infer(body, c, v, Gamma, ctx);
		ctx.pop();
		
		/* 
		 * Any occurence of 'var' in bodyType is equal to it's binder
		 * this is true from the way we construct the AST.
		 */
		return new ForallType(var, bodyType);
	}
	
	public TypeExpression inferLambdaApplication(LambdaApplication lapp, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		/**
		 *  The checking of the two expressions must be independent of one another
		 *  This is the reason why we use two different constraints graphs 
		 */
		ConstraintGraph c1 = new ConstraintGraph().addAll(c);
		ConstraintGraph c2 = new ConstraintGraph().addAll(c);

		TypeExpression tFun = infer(lapp.getFun(),c1,v,Gamma, ctx);
		if( !(tFun instanceof FunctionType)){
			throw new TypeCheckingException("A non-function term is being applied: " + lapp.getFun().preetyPrint());
		}
		TypeExpression src = ((FunctionType)tFun).getSrc();
		TypeExpression dst = ((FunctionType)tFun).getDest();
		
		check(lapp.getArg(), c2, v, Gamma, src, ctx, false);
		
		c.clear();
		c.addAll(c1);
		c.addAll(c2);
		return dst;
	}
	
	public TypeExpression inferTypeApplication(TypeApplication tapp, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		TypeExpression tParam = tapp.getArg();
		annot(v, tParam);
		
		TypeExpression tExp = infer(tapp.getFun(), c, v, Gamma, ctx);
		if( !(tExp instanceof ForallType)){
			throw new TypeCheckingException("A non-type-abstraction term is being applied : " + tapp.getFun());
		}
		
		TypeExpression bodyType = ((ForallType)tExp).getBody();
		TypeVariable var = ((ForallType)tExp).getQuant();
		
		return subst(bodyType, var, tParam);
	}

	public TypeExpression inferDataConstructor(DataConstructor cons, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		TypeExpression tExp = Gamma.get(cons);
		/* 
		 * Get a fresh copy of the type of the constructor 
		 * and annotate it with a fresh inference stage variable 
		 */
		tExp = tExp.freshCopy(new Context<TypeVariable>());
		
		int alpha = Stage.freshStageVar(v);
		v.add(alpha);
		Stage alphaS = Stage.inferenceStageVariable(alpha);
		instDCType(tExp, alphaS );
		
		return tExp;
	}

	public TypeExpression inferCaseAnalisys(CaseAnalisys ca, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		TypeExpression caseType = ca.getType();
		annot(v, caseType);
		
		/* 
		 * Reserve a variable for the annotation of the 
		 * case-analised expression's constructor 
		 */
		int alpha = Stage.freshStageVar(v);
		v.add(alpha);
		Stage alphaS = Stage.inferenceStageVariable(alpha);
		
		/* The type of the analyzed expression should be a type constructor application */
		TypeExpression eType = infer(ca.getExp(), c, v, Gamma, ctx);
		if(!(eType instanceof TypeConstructorApplication)){
			throw new TypeCheckingException("A value that is not member of any datatype is case analised : " + ca.getExp());
		}

		/* ci are all the constructors of the datatype on which case analisys is performed */
		Set<DataConstructor> declaredCons = ((TypeConstructorApplication)eType).getCons().getIdt().keySet();
		Set<TermVariable> branchCons = ca.getBranches().keySet();
		if( !declaredCons.equals(branchCons)){
			throw new TypeCheckingException("Wrong constructors used in case analisys. Used : " + branchCons + " declared " + declaredCons);
		}
		
		
		/* The stage of the type of 'e' */
		TypeConstructorApplication caseTypeTCA = (TypeConstructorApplication) eType;
		Stage s = caseTypeTCA.getCons().getStage();
		
		for( Map.Entry<TermVariable, Expression> br : ca.getBranches().entrySet()){
			DataConstructor c_i = (DataConstructor) br.getKey();
			Expression e_i = br.getValue();
			TypeExpression desired = null;
			/* 
			 * Instantiate type params, and the stage variable
			 * in the constructor's type. And check if the expression
			 * on the current branch has the desired type.
			 */
			try{
				desired = inst(c_i, alphaS, caseTypeTCA.getParams(), caseType);
			} catch(TypeCheckingException tex){
				System.err.println("Error while checking case expression : " + ca.preetyPrint());
				throw tex;
			}
			check(e_i, c, v, Gamma, desired, ctx, false); 
		}
		
		
		/* Add the stage-decreasing constraint to the constraint graph */
		c.addConstraint(s, alphaS.succ());
		return caseType;
	}
	
	public TypeExpression inferLetrec(Letrec lr, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		Set<Integer> vStar = new HashSet<Integer>(), v1 = null;
		TypeExpression letrecType = lr.getType();

		/* Check the structure of the letrec type*/
		if(!(letrecType  instanceof FunctionType)){
			throw new TypeCheckingException("Recursive definition "+ lr.preetyPrint() + " is not a function");
		}
		FunctionType letrecTypeF = (FunctionType) letrecType;
		if(!(letrecTypeF.getSrc() instanceof TypeConstructorApplication)){
			throw new TypeCheckingException("Letrec can be used only to code recursion on a recursive datatype : " + lr.preetyPrint());			
		}
		TypeConstructorApplication letrecArgType = (TypeConstructorApplication) letrecTypeF.getSrc();	
		TypeExpression letrecResultType  = letrecTypeF.getDest();
		
		
		/* Annotate with stages the type of letrec expression, the argument must be starred */
		letrecArgType.getCons().setStarred();
		try{
			 v1 = annotrec(v, vStar, letrecType);
		} catch(TypeCheckingException tex){
			System.err.println("Error while checking recursive definition : " + lr.preetyPrint() );
			throw tex;
		}
		
		/* 
		 * The stage variable used to annotate the type 
		 * constructor of the value on which recursion 
		 * is performed 
		 */
		Stage alpha = letrecArgType.getCons().getStage();
		
		/* 
		 * Check if every variable in starred position occur positively in
		 * the letrec type 
		 */
		// Keep the same free variables as in the original type since this type is used only for subtiping
		FunctionType letrecTypeSucc = (FunctionType) letrecType.freshCopy(ctx);
		letrecTypeSucc = (FunctionType) TypeSucc.succ(letrecTypeSucc, vStar);
		SubtypingConstraintGenerator.cgen(letrecResultType, letrecTypeSucc.getDest(), c, ctx);
		
		
		/** 
		 * Check the body first with types in ID.
		 * IF 'c' does not have any NT <> ID constraints, THEN continue with recCheck.
		 * ELSE retry with NT.
		 * IF recCheck fails THEN we are in the case letrec f=f => try NT
		 */
		
		/* save the current constraint graph */
		ConstraintGraph ctmp = new ConstraintGraph().addAll(c);
		int err = InferType.errors;
		
		Gamma.put(lr.getVar(), lr.getType());
		try{
			printEmph("Trying ID type for recursive definition of \"" + lr.getVar()+"\"");
			/* Check the body with types in ID */
			check( lr.getExp(), c, v, Gamma, letrecTypeSucc, ctx, true);
		}catch(TypeCheckingException tex){

			/* if ID <> NT constraints were added */
			if(tex.isNtConstraint()){
				/* restore the previous state */
				Gamma.remove(lr.getVar());
				c = ctmp;
				InferType.errors = err;
				
				/* try checking with types in NT */
				printEmph("FAILED : Retrying with NT type");
				NTAnnot.annot(lr.getType());
				NTAnnot.annot(letrecTypeSucc);
				Gamma.put(lr.getVar(), lr.getType());
				check( lr.getExp(), c, v, Gamma, letrecTypeSucc, ctx, false);						
				Gamma.remove(lr.getVar());
				printEmph("SUCCEEDED! " + lr.getType());
				return lr.getType();
			}else{
				throw tex;
			}
		}
		Gamma.remove(lr.getVar());

		/* Invoke the recCheck algorithm */
		try{
			c.recCheck(alpha, vStar, v1);
		}catch(TypeCheckingException tce){
			System.err.println("Recursive definition \n" + lr.preetyPrint() + "\nis believed to be non terminating");

			if( lr.getExp().equals(lr.getVar())){
				printEmph("FAILED : Retrying with NT type");
				NTAnnot.annot(lr.getType());
			} else {
				throw tce;
			}
		}
		printEmph("SUCCEEDED! " + letrecType.preetyPrint());		
		return letrecType;
	}
	
	private void printEmph(String msg){
		int len = msg.length() + 4;
		byte [] stars = new byte[len];
		Arrays.fill(stars, (byte)'-');
		System.err.println(new String(stars));
		System.err.println("| " + msg + " |");
		System.err.println(new String(stars));
	}
	
	public TypeExpression inferBindUnit(String name, LambdaApplication la, ConstraintGraph c, Set<Integer> v,TypingContext Gamma, Context<TypeVariable> ctx ){
		Expression fn = la.getFun();
		/* check wether it is a bind construction */
		if(!(fn instanceof TermVariable)){
			return null;
		} 
		TermVariable fnv = (TermVariable) fn;
		if(!fnv.getName().equals(name)){
			return null;
		}
		
		Expression arg = la.getArg();
		TypeExpression te = infer(arg, c, v, Gamma, ctx);
		
		/* for _unit_ annotate the entire type with the NT variable */
		if(name.equals("unit")){
			NTAnnot.annot(te);
		}
		
		/* for bind annot only the source argument */
		if(name.equals("bind")){
			if(!(te instanceof FunctionType)){
				throw new TypeCheckingException("bind can be applied only to a function. Found it applied to : " + la.getArg());
			}
			FunctionType fte = (FunctionType) te;			
			NTAnnot.annot(fte.getSrc());
		}
		return te;
	}
	
	/**
	 * The type checking function.
	 * @param exp	The expression
	 * @param v		The set of stage variables used to annotate types so far 
	 * @param Gamma The typing environment
	 * @param type	The type to check against
	 * @param ctx 	The context of free type variables
	 * @param rethrow TODO
	 */
	public void check(Expression exp, ConstraintGraph c, Set<Integer> v, TypingContext Gamma, TypeExpression type, Context<TypeVariable> ctx, boolean rethrow){
		TypeExpression infrd = infer(exp, c, v, Gamma, ctx);
		/* The desired type should be a subtype of the inferred one */
		try{
			SubtypingConstraintGenerator.cgen(infrd, type, c, ctx);
		}catch(TypeCheckingException te){
			/* re-throw the exception to the surrounding letrec rule */
			if( rethrow) 
				throw te;
			InferType.errors++;
			System.err.println("Subtyping error");
			System.err.println(exp.preetyPrint());
			System.err.println("should be of type ");
			System.err.println(type.preetyPrint());
			System.err.println("but found of type ");
			System.err.println(infrd.preetyPrint());
			System.err.println("==============================================");
		}
	}
	
	public TypeExpression infer(Expression exp, ConstraintGraph c, Set<Integer> v, TypingContext Gamma, Context<TypeVariable> ctx){
		TypeExpression ret = null;

		if( exp instanceof LambdaAbstraction){
			ret = inferLambdaAbstraction((LambdaAbstraction)exp, c, v, Gamma, ctx);
		} else if (exp instanceof LambdaApplication){
			/* check first if there is a bind or unit contruction */
			if( (ret = inferBindUnit("bind",(LambdaApplication) exp, c, v, Gamma, ctx)) == null){
				if( (ret = inferBindUnit("unit",(LambdaApplication) exp, c, v, Gamma, ctx)) == null){
					ret = inferLambdaApplication((LambdaApplication) exp, c, v, Gamma, ctx);							
				}
			}
		} else if (exp instanceof TypeAbstraction){
			ret = inferTypeAbstraction((TypeAbstraction) exp, c, v, Gamma, ctx);
		} else if (exp instanceof TypeApplication){
			ret = inferTypeApplication((TypeApplication) exp, c, v, Gamma, ctx);
		} else if (exp instanceof CaseAnalisys){
			ret = inferCaseAnalisys((CaseAnalisys) exp, c, v, Gamma, ctx);
		} else if (exp instanceof Letrec){
			ret = inferLetrec((Letrec) exp, c, v, Gamma, ctx);
		} else if (exp instanceof DataConstructor){
			ret = inferDataConstructor((DataConstructor)exp , c, v, Gamma, ctx);
		} else if (exp instanceof TermVariable ){
			ret = inferVariable((TermVariable) exp, c, v, Gamma, ctx);
		}
		c.simplify();
//		System.out.println("infered " + exp.getClass() + " {"  + exp.preetyPrint() + "} => " + ret.preetyPrint()) ;
		return ret;
	}

	/**
	 * Replaces the variable 'var', with the expression 'replacement' in the expression 'exp'.
	 */
	public TypeExpression subst(TypeExpression exp, TypeVariable var,  TypeExpression replacement){
		return TypeSubst.subst(exp, var, replacement);
	}

	/**
	 * Annotates the type expression 'type' with fresh stage variables
	 * with respect to the set 'v'.
	 */
	public void annot(Set<Integer> v, TypeExpression type){
		TypeAnnot.annot(type, v, null);
	}
	
	/**
	 * Annotates the type expression 'type' with fresh stage variables wrt
	 * v, vStar and puts all variables appearing in starred positions 
	 * in the vStar set. 
	 * 
	 * @param type The bare type of the letrec expression
	 */
	public Set<Integer> annotrec(Set<Integer> v, Set<Integer> vStar, TypeExpression type){
		TypeAnnot.annot(type, v, vStar);
		return new HashSet<Integer>(v);
	}

	/**
	 * Instantiates the type of a data constructor with the 
	 * variable 'alpha'. 
	 */
	public void instDCType(TypeExpression te , Stage alpha){
		InstConstructorType.inst(te, alpha);
	}
		
	public TypeExpression inst(DataConstructor c, Stage alpha, List<TypeExpression> tParams, TypeExpression retType){
		TypeExpression consType = c.getIdt().get(c);
		/* get a fresh copy of the type that is to instantiate with given variable */
		consType = consType.freshCopy(new Context<TypeVariable>());
		instDCType(consType, alpha);
		
		/* 
		 * Apply the type of the constructor to it's arguments 
		 * and check if the exact number of arguments is supplied.
		 * (kind checking)
		 */
		for( int i = 0 ; i < tParams.size(); i++){
			if(! (consType instanceof ForallType)){
				throw new TypeCheckingException("Too many arguments given to type constructor.");
			}
			ForallType consTypeFA = (ForallType) consType;
			consType = subst(consTypeFA.getBody(), consTypeFA.getQuant(), tParams.get(i));
		}

		if((consType instanceof ForallType)){
			throw new TypeCheckingException("Too few arguments given to type constructor.");
		}
		
		return replaceRightmost(consType, retType);
	}
	
	private TypeExpression replaceRightmost(TypeExpression te, TypeExpression ret){		
		if(te instanceof FunctionType){
			FunctionType tefn = (FunctionType) te;
			return new FunctionType(tefn.getSrc(), replaceRightmost(tefn.getDest(), ret));
		} else {
			return ret;
		}
	}
	
}
